import json
import threading
import airbot
import time
import os
from utils_py import utils_set_zero
from utils_py.utils_motor_head import Head
from utils_py.utils_positions import *
from utils_py.utils_limx_ws import limx_ws
import utils_py.utils_gamepad_limx as limx_gamepad
from utils_py.configure import MotorCfg
from utils_py.utils_motor_waist import Waist
from functools import partial
from utils_py.utils_hands import *

# 获取当前文件的绝对路径
current_file_path = os.path.abspath(__file__)
# 获取当前文件所在的文件夹路径
current_dir = os.path.dirname(current_file_path)

class Human:

    def __init__(self, ACCID="WF_TRON1A_086", leg_ws_url="ws://10.192.1.2:5000", left_arm_can="can0", right_arm_can="can1", open_notify=False):
        """Human 类初始化"""
        self.head = Head('/dev/ttyUSB0')
        self.should_exit = False
        self.left_arm = None
        self.right_arm = None
        self.left_arm_can = left_arm_can
        self.right_arm_can = right_arm_can
        self.leg_ws = limx_ws(ACCID, leg_ws_url, open_notify)
        self.waist = Waist()
        # self.hands = RH56HandController()
    def control_hands(self,r1,r2,r3,r4,r5,r6,l1,l2,l3,l4,l5,l6):
        controller1 = self.RH56HandController(port="/dev/ttyUSB1")  # 修改为实际串口
        controller2 = self.RH56HandController(port="/dev/ttyUSB2")  # 修改为实际串口
        # try:
            # while True:
                # 读取实际角度
        actual_angles1 = controller1.get_actual_angles()
        print(f"1实际角度值: {actual_angles1}")
        actual_angles2 = controller2.get_actual_angles()
        # print(f"2实际角度值: {actual_angles2}")
        # 设置目标角度（示例：设置各角度为100,100,100,100,2000,0）
        # target_angles = [500, 950, 950, 950, 950, 950]
        # target_angles = [r1,r2,r3,r4,r5,r6]
        # target_angles = []
        controller1.set_target_angles([r1,r2,r3,r4,r5,r6])
        controller2.set_target_angles([l1,l2,l3,l4,l5,l6])
        print("目标角度已设置")
        time.sleep(2)
        # target_angles = [200, 200, 200, 200, 500, 800]
        # controller1.set_target_angles(target_angles)
        # controller2.set_target_angles(target_angles)
        # print("目标角度已设置")
        # time.sleep(2)
    def connect_left_arm(self):
        """连接左臂"""
        print("Connecting left arm...")
        self.left_arm = None
        time.sleep(1)
        self.left_arm = airbot.create_agent("down", self.left_arm_can, 0.7, "none")
        self.left_arm.online_mode()
        time.sleep(2)

    def connect_right_arm(self):
        """连接右臂"""
        print("Connecting right arm...")
        self.right_arm = None
        time.sleep(1)
        self.right_arm = airbot.create_agent("down", self.right_arm_can, 0.7, "none")
        self.right_arm.online_mode()
        time.sleep(2)

    def left_arm_go_home(self):
        """初始化左手"""
        try:
            print("Left arm go home...")
            assert self.left_arm is not None, "left arm not connected, please check out"
            ret = self.left_arm.set_target_pose(left_arm_home)
            print(f"left_arm ret {ret}")
        except Exception as e:
            print(f"Left arm go home error: {e}")
            print("Please checkout")

    def right_arm_go_home(self):
        """初始化右手"""
        try:
            print("Right arm go home...")
            assert self.right_arm is not None, "right arm not connected, please check out"
            ret = self.right_arm.set_target_pose(right_arm_home)
            print(f"right_arm ret {ret}")
        except Exception as e:
            print(f"Right arm go home error: {e}")
            print("Please checkout")

    def arm_go_dance1(self):
        """初始化右手"""
        try:
            # zuoshou
            interval = 1.0
            print("left_arm_go_dance...")
            assert self.left_arm is not None, "right arm not connected, please check out"
            ret=self.left_arm.set_target_pose(left_arm_dance1)
            print(f"left_arm ret {ret}")
            time.sleep(interval)
            self.left_arm.set_target_pose(left_arm_home)
            time.sleep(interval)

            # youshou
            ret=self.right_arm.set_target_pose(right_arm_dance1)
            print(f"left_arm ret {ret}")
            time.sleep(interval)
            self.right_arm.set_target_pose(right_arm_home)
            time.sleep(interval)
            # youshou
            ret=self.right_arm.set_target_pose(right_arm_dance1)
            print(f"left_arm ret {ret}")
            time.sleep(interval)
            self.right_arm.set_target_pose(right_arm_home)
            time.sleep(interval)
            # zuoshou
            print("left_arm_go_dance...")
            assert self.left_arm is not None, "right arm not connected, please check out"
            ret=self.left_arm.set_target_pose(left_arm_dance1)
            print(f"left_arm ret {ret}")
            time.sleep(interval)
            self.left_arm.set_target_pose(left_arm_home)
            time.sleep(interval)
            # yiqi
            ret=self.left_arm.set_target_pose(left_arm_dance1)
            ret=self.right_arm.set_target_pose(right_arm_dance1)
            time.sleep(2)
            self.left_arm.set_target_pose(left_arm_home)
            self.right_arm.set_target_pose(right_arm_home)

        except Exception as e:
            print(f"left_arm_go_danceerror:{e}")
            print("Please checkout")
    
    def arm_go_dance2(self):
        try:
            # zuoshou
            interval = 1.0
            # print("left_arm_go_dance...")
            # assert self.left_arm is not None, "right arm not connected, please check out"
            # ret=self.left_arm.set_target_pose(left_arm_dance2)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval)
            # self.left_arm.set_target_pose(left_arm_home)
            # time.sleep(interval)

            # # youshou
            # ret=self.right_arm.set_target_pose(right_arm_dance2)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval)
            # self.right_arm.set_target_pose(right_arm_home)
            # time.sleep(interval)
            # # youshou
            # ret=self.right_arm.set_target_pose(right_arm_dance2)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval)
            # self.right_arm.set_target_pose(right_arm_home)
            # time.sleep(interval)
            # # zuoshou
            # print("left_arm_go_dance...")
            # assert self.left_arm is not None, "right arm not connected, please check out"
            # ret=self.left_arm.set_target_pose(left_arm_dance2)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval)
            # self.left_arm.set_target_pose(left_arm_home)
            # time.sleep(interval)
            # yiqi
            ret=self.left_arm.set_target_pose(left_arm_dance2)
            ret=self.right_arm.set_target_pose(right_arm_dance2)
            time.sleep(2)
            self.left_arm.set_target_pose(left_arm_home)
            self.right_arm.set_target_pose(right_arm_home)

        except Exception as e:
            print(f"left_arm_go_danceerror:{e}")
            print("Please checkout")
    
    def arm_go_dance3(self):
        try:
            # zuoshou
            interval = 1.0
            # print("left_arm_go_dance...")
            # assert self.left_arm is not None, "right arm not connected, please check out"
            # ret=self.left_arm.set_target_pose(left_arm_dance31)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval)
            # self.left_arm.set_target_pose(left_arm_home)
            # time.sleep(interval)

            # # youshou
            # ret=self.right_arm.set_target_pose(right_arm_dance32)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval)
            # self.right_arm.set_target_pose(right_arm_home)
            # time.sleep(interval)
            # # youshou
            # ret=self.right_arm.set_target_pose(right_arm_dance31)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval)
            # self.right_arm.set_target_pose(right_arm_home)
            # time.sleep(interval)
            # # zuoshou
            # print("left_arm_go_dance...")
            # assert self.left_arm is not None, "right arm not connected, please check out"
            # ret=self.left_arm.set_target_pose(left_arm_dance32)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval)
            # self.left_arm.set_target_pose(left_arm_home)
            # time.sleep(interval)
            
            # yiqi
            # zuoshou
            interval = 1.0
            print("left_arm_go_dance...")
            assert self.left_arm is not None, "right arm not connected, please check out"
            ret=self.left_arm.set_target_pose(left_arm_dance31)
            ret=self.right_arm.set_target_pose(right_arm_dance32)
            print(f"left_arm ret {ret}")
            time.sleep(interval*2)
            ret=self.right_arm.set_target_pose(right_arm_dance31)
            ret=self.left_arm.set_target_pose(left_arm_dance32)
            print(f"left_arm ret {ret}")
            time.sleep(interval*2)

            # ret=self.left_arm.set_target_pose(left_arm_dance31)
            # ret=self.right_arm.set_target_pose(right_arm_dance32)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval*2)
            # ret=self.right_arm.set_target_pose(right_arm_dance31)
            # ret=self.left_arm.set_target_pose(left_arm_dance32)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval*2)

            # ret=self.left_arm.set_target_pose(left_arm_dance31)
            # ret=self.right_arm.set_target_pose(right_arm_dance32)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval*2)
            # ret=self.right_arm.set_target_pose(right_arm_dance31)
            # ret=self.left_arm.set_target_pose(left_arm_dance32)
            # print(f"left_arm ret {ret}")
            # time.sleep(interval*2)
            
            self.right_arm.set_target_pose(right_arm_home)
            self.left_arm.set_target_pose(left_arm_home)
            time.sleep(interval)

        except Exception as e:
            print(f"left_arm_go_danceerror:{e}")
            print("Please checkout")
    
    def arm_go_dance4(self):
        try:
            # zuoshou
            interval = 1.5
            print("left_arm_go_dance...")
            assert self.left_arm is not None, "right arm not connected, please check out"
            ret=self.left_arm.set_target_pose(left_arm_dance41)
            print(f"left_arm ret {ret}")
            time.sleep(interval)
            self.left_arm.set_target_pose(left_arm_dance42)
            time.sleep(interval)
            self.left_arm.set_target_pose(left_arm_dance41)
    

            # youshou
            ret=self.right_arm.set_target_pose(right_arm_dance41)
            print(f"left_arm ret {ret}")
            time.sleep(interval)
            self.right_arm.set_target_pose(right_arm_dance42)
            time.sleep(interval)
            self.right_arm.set_target_pose(right_arm_dance41)
            # youshou
            ret=self.right_arm.set_target_pose(right_arm_dance41)
            print(f"left_arm ret {ret}")
            time.sleep(interval)
            self.right_arm.set_target_pose(right_arm_dance42)
            time.sleep(interval)
            self.right_arm.set_target_pose(right_arm_dance41)
            # zuoshou
            ret=self.left_arm.set_target_pose(left_arm_dance41)
            print(f"left_arm ret {ret}")
            time.sleep(interval)
            self.left_arm.set_target_pose(left_arm_dance42)
            time.sleep(interval)
            self.left_arm.set_target_pose(left_arm_dance41)
            # yiqi
            ret=self.left_arm.set_target_pose(left_arm_dance41)
            ret=self.right_arm.set_target_pose(right_arm_dance41)
            time.sleep(2)
            self.left_arm.set_target_pose(left_arm_dance42)
            self.right_arm.set_target_pose(right_arm_dance42)
            time.sleep(2)
            ret=self.left_arm.set_target_pose(left_arm_dance41)
            ret=self.right_arm.set_target_pose(right_arm_dance41)

        except Exception as e:
            print(f"left_arm_go_danceerror:{e}")
            print("Please checkout")

    def go_zero(self):
        """左右手回到零点位置"""
        try:
            print("Go zero position.")
            ret1=self.left_arm.set_target_pose([0.14362652818848726, -0.0048399906702387575, 0.2149357327769668], [-0.003598171043691348, 0.006618258168461989, -0.028393032405924036, 0.9995684506499553],use_planning=True)
            print(f"left_arm ret {ret1}")
            ret=self.right_arm.set_target_pose([0.14362652818848726, -0.0048399906702387575, 0.2149357327769668], [-0.003598171043691348, 0.006618258168461989, -0.028393032405924036, 0.9995684506499553],use_planning=True)
            print(f"right_arm ret {ret}")
            time.sleep(5) # 等待一段时间，确保手回到零点位置
        except Exception as e:
            print(f"Go zero position  error:{e}")
            print("Please checkout")
    
    def go_home(self):
        """左右手回到初始位置（双手自然下垂）"""
        try:
            print("Go home position.")
            ret1=self.left_arm.set_target_pose(left_arm_home)
            print(f"left_arm ret {ret1}")
            ret=self.right_arm.set_target_pose(right_arm_home)
            print(f"right_arm ret {ret}")
            time.sleep(5) # 等待一段时间，确保手回到初始位置
        except Exception as e:
            print(f"Go home position error:{e}")
            print("Please checkout")

    def replay(self, actionName):
        """回放动作"""
        def __play_actions(records_json_path, robot):
            with open(records_json_path) as records_json:
                traj_list = json.load(records_json)
            # 控制机器人
            for traj in traj_list:
                robot.set_target_joint_q(traj, use_planning=False)
                # robot.set_target_joint_mit(traj,[0, 0, 0, 0, 0, 0],[120, 120, 120, 30, 30, 30],[1.2, 1.2, 1.2, 0.15, 0.15, 0.15],)
                ret = robot.get_current_pose()
                # print(ret)
                time.sleep(0.001)
            time.sleep(2)
        # home point
            # ret=self.left_arm.set_target_pose(left_arm_home)
            # ret=self.right_arm.set_target_pose(right_arm_home)
    
            # time.sleep(5)
            # self.left_arm.manual_mode()
            # self.right_arm.manual_mode()
            # time.sleep(20)
        # ([0.14362652818848726, -0.0048399906702387575, 0.2149357327769668], [-0.003598171043691348, 0.006618258168461989, -0.028393032405924036, 0.9995684506499553])
        left_file = current_dir + '/action_left/' + actionName + '.json'
        right_file = current_dir + '/action_right/' + actionName + '.json'
        thread_left = None
        thread_right = None
        if os.path.exists(left_file) and self.left_arm is not None:
            thread_left = threading.Thread(target=__play_actions, args=(left_file, self.left_arm))
            thread_left.start()
            # thread_left.daemon = True

        if os.path.exists(right_file) and self.right_arm is not None:
            thread_right = threading.Thread(target=__play_actions, args=(right_file, self.right_arm))
            thread_right.start()
            # thread_right.daemon = True
        if thread_left is not None:
            thread_left.join()
        if thread_right is not None:    
            thread_right.join()   
    
    def calibrate_left_arm(self):
        """左手标定"""
        try:
            print("Calibrating left arm...")
            self.left_arm = None 
            time.sleep(2)
            #thread = threading.Thread(target=utils_set_zero.execute_command_with_enter, args=(f"airbot_set_zero -m {self.left_arm_can}",))
            #thread.start()
            utils_set_zero.execute_command_with_enter(f"airbot_set_zero -m {self.left_arm_can}")
            #thread.join()
            time.sleep(3)
            self.connect_left_arm() 
        except Exception as e:
            print(f"Calibrating left arm error:{e}")
            print("Please checkout")

    def calibrate_right_arm(self):
        """右手标定"""
        try:
            print("Calibrating right arm...")
            self.right_arm = None
            time.sleep(2)
            #thread = threading.Thread(target=utils_set_zero.execute_command_with_enter, args=(f"airbot_set_zero -m {self.right_arm_can}",))
            #thread.start()
            #thread.join()
            utils_set_zero.execute_command_with_enter(f"airbot_set_zero -m {self.right_arm_can}")
            time.sleep(3)
            self.connect_right_arm()
        except Exception as e:
            print(f"Calibrating right arm error:{e}")
            print("Please checkout")

# region 原来的跳舞程序
def demo():
    # 导入游戏手柄相关的工具函数
    from utils_py import utils_gamepad

    human = Human(ACCID="WF_TRON1A_086")                                                                       
    # human.open_notify = True # 是否打印实时反馈日志 默认False
    # human.start() #
    # time.sleep(2)

    # 提示用户需要先校准双臂
    print("Step 1: Please calibrate both arms first.")

    # 无限循环，等待游戏手柄的输入
    while(True):
        # 提示用户正在等待游戏手柄输入
        print("Waiting for gamepad input...")

        # 获取游戏手柄的操作动作
        action = utils_gamepad.get_gamepad_action2()

        # 打印接收到的动作信息
        print(f"Received action: {action}")

        # 如果没有接收到有效的动作，提示该动作没有配置
        if action is None:
            print("Action is not configured. Please try again.")
            continue

        # 根据不同的动作执行相应的操作
        if action == "L_Zero":
            # 左臂标零，回归到初始位置（双臂自然下垂）
            human.calibrate_left_arm()
            if(human.left_arm is None):
                print("Left arm connection failed, please checkout")
                continue
            human.left_arm_go_home()
        
        elif action == "R_Zero":
            # 将右臂标零，回归到初始位置（双臂自然下垂）
            human.calibrate_right_arm()
            if(human.right_arm is None):
                print("Right arm connection failed, please checkout")
                continue
            human.right_arm_go_home()
        
        elif action == "go_zero":
            human.go_zero() # 双臂回到零位
            
        elif action == "go_home":
            human.go_home() # 将双臂移动到初始位置（双臂自然下垂）
        elif action == "Dance1":
            # 左臂标零，回归到初始位置（双臂自然下垂）
            human.arm_go_dance1()
        elif action == "Dance2":
            human.arm_go_dance2()
        elif action == "Dance3":
            # 左臂标零，回归到初始位置（双臂自然下垂）
            human.arm_go_dance3()
        elif action == "Dance4":
            human.arm_go_dance4()
        elif action == "Danceall":
            human.arm_go_dance1()
            time.sleep(2)
            human.arm_go_dance2()
            time.sleep(1)
            human.arm_go_dance3()
            time.sleep(1)
            human.arm_go_dance1()
            human.go_zero()
            human.arm_go_dance4()
            human.replay('wave')
            # human.replay('say_goodbye')
            human.go_zero()
          
        else:
            # 对于其他动作，进行录播操作
            human.replay(action)
            print(f"Replaying action: {action}")
            # 等待5秒钟
            time.sleep(2)

#endregion

class JoyReceiver(limx_gamepad.LimxReceiver):
    def __init__(self):
        super().__init__()

    def play_action(self, action):
        super().play_action(action)
        
        if action is None or action == "None":
            print("Action is not configured. Please try again.")
            return 

        elif action == "head_left":
            human.head.left(360)
            return
        elif action == "head_right":
            human.head.right(360)
            return
        elif action == "head":
            human.head.left_right()
            return
        
        elif action == "waist_up":
            human.waist.up()
            return
        elif action == "waist_down":
            human.waist.down()
            return
        elif action == "waist_left":
            human.waist.left()
            return
        elif action == "waist_right":
            human.waist.right()
        elif action == "hand_grip":
            human.control_hands([500, 950, 950, 950, 950, 950,500, 950, 950, 950, 950, 950])
            return
        elif action == "hand_full":
            human.control_hands([200, 200, 200, 200, 500, 800,200, 200, 200, 200, 500, 800])
            return
        # 根据不同的动作执行相应的操作
        elif action == "L_Zero":
            # 左臂标零，回归到初始位置（双臂自然下垂）
            human.calibrate_left_arm()
            if(human.left_arm is None):
                print("Left arm connection failed, please checkout")
                return
            #human.left_arm_go_home()
        
        elif action == "R_Zero":
            # 将右臂标零，回归到初始位置（双臂自然下垂）
            human.calibrate_right_arm()
            if(human.right_arm is None):
                print("Right arm connection failed, please checkout")
                return
            #human.right_arm_go_home()
        
        elif action == "go_zero":
            human.go_zero() # 双臂回到零位
            
        elif action == "go_home":
            human.go_home() # 将双臂移动到初始位置（双臂自然下垂）
        elif action == "Dance1":
            # 左臂标零，回归到初始位置（双臂自然下垂）
            human.arm_go_dance1()
        elif action == "Dance2":
            human.arm_go_dance2()
        elif action == "Dance3":
            # 左臂标零，回归到初始位置（双臂自然下垂）
            human.arm_go_dance3()
        elif action == "Dance4":
            human.arm_go_dance4()
        elif action == "Danceall":
            human.arm_go_dance1()
            time.sleep(2)
            #human.arm_go_dance2()
            #time.sleep(1)
            human.arm_go_dance3()
            time.sleep(1)
            human.arm_go_dance1()
            human.go_zero()
            human.arm_go_dance4()
            human.replay('wave')
            # human.replay('say_goodbye')
            human.go_zero()
          
        else:
            # 对于其他动作，进行录播操作
            human.replay(action)
            print(f"Replaying action: {action}")
            # 等待5秒钟
            time.sleep(2)
    
human = Human()                                                                       

# human.start() #
# time.sleep(2)
# 提示用户需要先校准双臂
print("Step 1: Please calibrate both arms first.")

if __name__ == "__main__":
    robot = limx_gamepad.initialize_robot("10.192.1.2")

    # Create an instance of RobotReceiver to handle callbacks
    receiver = JoyReceiver()

    # Create partial functions for callbacks
    sensorJoyCallback = partial(receiver.sensorJoyCallback)

    robot.subscribeSensorJoy(sensorJoyCallback)
    
    while(True):
        time.sleep(0.01)



